﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Xml;
using System.IO;
using System.Diagnostics;

namespace XmlDSigEx
{
    public class Reference
    {
        private Dictionary<string, XmlElement> _transforms = new Dictionary<string, XmlElement>();
        private List<ITransform> _transformList = new List<ITransform>();
       
        internal Reference()
        {
            Uri = null;
            DigestMethod = CryptoConfigDigestAlgorithm.URI_DIGEST_SHA256;
        }

        public string Uri { get; set; }

        public string Type { get; set; }

        internal void AddTransform(string uri, XmlElement element) { _transforms.Add(uri, element); }

        public string DigestMethod { get; set; }

        public byte[] DigestValue { get; set; }

        public bool Precomputed { get; set; }


        internal void WriteXml(XmlWriter writer)
        {
            writer.WriteStartElement("Reference", XmlSigner.NS_XMLDSIG);
            if (Uri != null)
                writer.WriteAttributeString("URI", Uri);

            if (Type != null)
                writer.WriteAttributeString("Type", Type);      

            WriteTransforms(writer);

            writer.WriteStartElement("DigestMethod", XmlSigner.NS_XMLDSIG);
            writer.WriteAttributeString("Algorithm", DigestMethod);
            writer.WriteEndElement();

            writer.WriteStartElement("DigestValue", XmlSigner.NS_XMLDSIG);
            if (Precomputed)
                writer.WriteString(Convert.ToBase64String(DigestValue));
            writer.WriteEndElement();
            writer.WriteEndElement();
        }

        private void WriteTransforms( XmlWriter writer)
        {
            if (_transformList.Count > 0)
            {
                writer.WriteStartElement("Transforms", XmlSigner.NS_XMLDSIG);

                foreach (ITransform transform in _transformList)
                {
                    writer.WriteStartElement("Transform", XmlSigner.NS_XMLDSIG);
                    writer.WriteAttributeString("Algorithm", transform.Uri);
                    transform.WriteXml(writer);
                    writer.WriteEndElement();
                }
                writer.WriteEndElement();
            }
        }

        internal void ProcessDigest(XmlDocument context, XmlElement reference)
        {
            XmlNamespaceManager nsMgr = new XmlNamespaceManager(reference.OwnerDocument.NameTable);
            nsMgr.AddNamespace("ds", XmlSigner.NS_XMLDSIG);
            byte[] digest = CalculateDigest(context, reference);
            reference.SelectSingleNode("./ds:DigestValue", nsMgr)
                .InnerText = Convert.ToBase64String(digest);
        }

        private byte[] CalculateDigest(XmlDocument context, XmlElement referenceElement)
        {
            IUriDereferencer dereferencer = DereferencerFactory.GetDereferencer(Uri);
            object input = dereferencer.DereferenceUri(Uri, context);            
            Stream output = null;
            byte[] digest = null;
            try
            {
                if (_transforms.Count > 0)
                {
                    output = ProcessTransforms(input, referenceElement, context);
                }

                if (output == null)
                {
                    output = CanonicaliseNodeSet(input);
                }



                IDigestAlgorithm digestAlgorithm = DigestFactory.GetDigestAlgorithm(DigestMethod);
                digest = digestAlgorithm.CreateDigest(output);
            }
            finally
            {
                IDisposable stream = input as IDisposable;
                if (stream != null)
                    stream.Dispose();
            }
            return digest;
        }

        private Stream CanonicaliseNodeSet(object input)
        {
            List<XmlNode> nodeSet = input as List<XmlNode>;
            if (nodeSet == null)
                throw new ArgumentException("CanonicaliseNodeSet must be passed a List<XmlNode>.");

            ICanonicalisationAlgorithm canonicalisation = new XmlC14NCanonicalisation();
            return canonicalisation.ProcessNodeSet(nodeSet, false);
        }

        private Stream ProcessTransforms(object input, XmlElement referenceElement, XmlDocument documentContext)
        {
            List<ITransform> transformObjects = new List<ITransform>();

            foreach (KeyValuePair<string, XmlElement> transformPair in _transforms)
            {
                ITransform transform = TransformFactory.GetTransform(transformPair.Key, documentContext);
                transformObjects.Add(transform);
            }

            TransformChain chain = new TransformChain(referenceElement, transformObjects.ToArray());
            return chain.ProcessTransforms(input);
        }

        internal bool Check(XmlDocument context, XmlElement referenceElement)
        {
            byte[] digest = CalculateDigest(context, referenceElement);
            return digest.IsEqualTo(DigestValue);
        }

        public ITransform AddTransform(string uri)
        {
            AddTransform(uri, null);
            try
            {
                ITransform transform = TransformFactory.GetTransform(uri, null);
                _transformList.Add(transform);
                return transform;
            }
            catch (UnknownAlgorithmException exc)
            {
                if (Precomputed)
                {
                    _transformList.Add(new NullTransform(uri));
                    return null;
                }
                else
                {
                    throw;
                }
            }
        }

        public IEnumerable<string> Transforms
        {
            get
            {
                return _transforms.Keys.Cast<string>();
            }
        }
    }
}
